﻿from typing import List, Dict
from typedef import *
from memory import GlobalMemory
from cache import Cache
from tlb import MMU
from dataclasses import dataclass, field
from enum import Enum
import numpy as np
from time import sleep
from threading import Thread
from numpy import ndarray, uint8, uint16, uint32, uint64, float32, float64


eq = 1
gt = 2
lt = 4
abs_addr = 8


@dataclass
class Command:
    opcode: uint8
    r1: uint8
    r2: uint8
    r3: uint8
    imm: int

    def __init__(self, value: int):
        self.opcode = value & 0x7f
        self.r1 = (value >> 7) & 0x1f
        self.r2 = (value >> 12) & 0x1f
        self.r3 = (value >> 17) & 0x1f
        self.imm = (value >> 22) & 0x3ff
        self.imm = self.imm if (self.imm < 0x200) else (self.imm-0x400)


class Core:
    __iregs: ndarray[uint64] = field(
        default_factory=lambda: np.zeros(shape=(2, 32), dtype=uint64))
    __fregs: ndarray[float32] = field(
        default_factory=lambda: np.zeros(shape=32, dtype=float32))
    __dregs: ndarray[float64] = field(
        default_factory=lambda: np.zeros(shape=32, dtype=float64))
    __sregs: List[List[int]] = [[0]*8]*2
    __ir: Command = Command(0)
    __mar: uint64 = 0
    __mdr: ndarray[uint8] = field(
        default_factory=lambda: np.zeros(shape=8, dtype=uint8))

    def __init__(self, id: int, memory: GlobalMemory, cache: Cache, tlb: MMU):
        self.id = id
        self.__cache = cache
        self.__tlb = tlb
        self.__memory = memory
        self.__iregs = np.zeros(shape=(2, 32), dtype=uint64)
        self.__mdr = np.zeros(shape=8, dtype=uint8)
        self.__iregs[0, Stack] = 0xf1000
        self.__iregs[1, Stack] = 0xff000
        self.__sregs[0][Status] = IntEN | CacheEnable | (
            Master if id == 0 else (Slave | Block))
        self.__sregs[1][Status] = IntEN | CacheEnable | (
            Master if id == 1 else (Slave | Block)) | Privilege
        self.__priv: int = 1
        self.__rwThread: Thread = Thread()

    def __fetch(self):
        self.__tlb.update(self.__sregs[self.__priv])
        self.__mar = self.__sregs[self.__priv][Code] << 2
        self.__rwThread = Thread(target=self.__ram, daemon=True,
                                 args=(4, Read | Execute | (Priv if self.__priv == 1 else 0)))
        self.__rwThread.start()
        self.__sregs[self.__priv][Code] += 1

    def __ram(self, size: int, access: int):
        self.__mar = self.__tlb(self.__mar, access)
        if access & Read:
            self.__mdr.fill(0)
            if size == 1:
                self.__mdr[0] = self.__cache.read(
                    self.__mar, size, access & Execute > 0)
            else:
                self.__mdr[0:size] = self.__cache.read(
                    self.__mar, size, access & Execute > 0)
        elif access & Write:
            self.__cache.write(self.__mar, self.__mdr, size, True)
        else:
            raise DefaultFault(
                self.__sregs[self.__priv][Code]-1, "Memory Operation Error")

    def __post(self, e: VMException):
        int_id = e.code
        self.__mar = self.__sregs[self.__priv][IDTR]
        self.__mar = self.__mar+(int_id << 3)
        self.__rwThread = Thread(target=self.__ram, daemon=True,
                                 args=(8, Read | (Priv if self.__priv else 0)))
        self.__rwThread.start()
        while self.__rwThread.is_alive():
            sleep(1e-7)
        self.__priv = 1
        self.__sregs[True][Code] = self.__mdr.view(dtype=uint64)[0]

    def __call__(self):
        while True:
            try:
                self.__fetch()
                while self.__rwThread.is_alive():
                    sleep(1e-7)
                self.__ir = Command(int(self.__mdr.view(dtype=uint32)[0]))
                match self.__ir.opcode:
                    case 0:  # nop
                        pass
                    case 1:  # calc
                        rd, rs, rt = self.__ir.r1, self.__ir.r2, self.__ir.r3
                        match self.__ir.imm:
                            case 1:  # move
                                self.__iregs[self.__priv][rd] = self.__iregs[self.__priv][rs]
                            case 2:  # add
                                self.__iregs[self.__priv][rd] = self.__iregs[self.__priv][rs] + \
                                    self.__iregs[self.__priv][rt]
                            case 4:  # sub
                                self.__iregs[self.__priv][rd] = self.__iregs[self.__priv][rs] - \
                                    self.__iregs[self.__priv][rt]
                            case 8:  # neg
                                self.__iregs[self.__priv][rd] = - \
                                    self.__iregs[self.__priv][rs]
                            case _:
                                raise IllegalInstruction(
                                    self.__sregs[self.__priv][Code]-1)
                    case 2:  # logic
                        rd, rs, rt = self.__ir.r1, self.__ir.r2, self.__ir.r3
                        match self.__ir.imm:
                            case 1:  # and
                                self.__iregs[self.__priv][rd] = self.__iregs[self.__priv][rs] & \
                                    self.__iregs[self.__priv][rt]
                            case 2:  # or
                                self.__iregs[self.__priv][rd] = self.__iregs[self.__priv][rs] | \
                                    self.__iregs[self.__priv][rt]
                            case 4:  # xor
                                self.__iregs[self.__priv][rd] = self.__iregs[self.__priv][rs] ^ \
                                    self.__iregs[self.__priv][rt]
                            case 8:  # not
                                self.__iregs[self.__priv][rd] = ~ \
                                    self.__iregs[self.__priv][rs]
                            case 16:  # shl
                                self.__iregs[self.__priv][rd] = self.__iregs[self.__priv][rs] << \
                                    rt
                            case 32:  # shr
                                self.__iregs[self.__priv][rd] = self.__iregs[self.__priv][rs] >> \
                                    rt
                            case 64:  # sar
                                self.__iregs[self.__priv][rd] = self.__iregs[self.__priv][rs] >> \
                                    rt
                            case _:
                                raise IllegalInstruction(
                                    self.__sregs[self.__priv][Code]-1)
                    case 3:  # set
                        rd, op, rt = self.__ir.r1, self.__ir.r2, self.__ir.r3
                        match op:
                            case 1:  # mov(unsigned)
                                self.__iregs[self.__priv][rd] = self.__ir.imm & 0x3ff
                            case 2:  # sign
                                self.__iregs[self.__priv][rd] = self.__ir.imm
                            case 4:  # add
                                self.__iregs[self.__priv][rd] = self.__ir.imm + \
                                    self.__iregs[self.__priv][rt]
                            case 8:  # sub
                                self.__iregs[self.__priv][rd] = self.__ir.imm - \
                                    self.__iregs[self.__priv][rt]
                            case 16:  # lsh
                                self.__iregs[self.__priv][rd] = (
                                    self.__ir.imm & 0x3ff) << rt
                            case 32:  # rsh
                                self.__iregs[self.__priv][rd] = (
                                    self.__ir.imm & 0x3ff) >> rt
                            case _:
                                raise IllegalInstruction(
                                    self.__sregs[self.__priv][Code]-1)
                    case 4:  # load
                        rd, ra, size = self.__ir.r1, self.__ir.r2, self.__ir.r3
                        match size:
                            case 1 | 2 | 4 | 8:
                                self.__mar = int(self.__iregs[self.__priv][ra]) + \
                                    self.__ir.imm
                                self.__tlb.update(self.__sregs[self.__priv])
                                self.__rwThread = Thread(target=self.__ram, daemon=True,
                                                         args=(size, Read | (Priv if self.__priv else 0)))
                                self.__rwThread.start()
                            case _:
                                raise AddrMisalign(
                                    self.__sregs[self.__priv][Code]-1)
                    case 5:  # store
                        rd, ra, size = self.__ir.r1, self.__ir.r2, self.__ir.r3
                        match size:
                            case 1 | 2 | 4 | 8:
                                self.__mar = int(self.__iregs[self.__priv][ra]) + \
                                    self.__ir.imm
                                self.__tlb.update(self.__sregs[self.__priv])
                                self.__mdr = self.__iregs[self.__priv][rd].view(
                                    dtype=uint8)
                                self.__rwThread = Thread(target=self.__ram, daemon=True,
                                                         args=(size, Write | (Priv if self.__priv else 0)))
                                self.__rwThread.start()
                            case _:
                                raise AddrMisalign(
                                    self.__sregs[self.__priv][Code]-1)
                    case 0x0a:  # jmp
                        r_cond, r_addr, cond = self.__ir.r1, self.__ir.r2, self.__ir.r3

                        def get_base(cond: int):
                            if cond & abs_addr:
                                return self.__iregs[self.__priv][r_addr]
                            else:
                                return self.__sregs[self.__priv][Code] & 0xffff_ffff_ffff_fc00
                        if cond & eq:
                            if self.__iregs[self.__priv][r_cond] == 0:
                                self.__sregs[self.__priv][Code] = get_base(
                                    cond)+self.__ir.imm
                        if cond & gt:
                            if self.__iregs[self.__priv][r_cond] > 0:
                                self.__sregs[self.__priv][Code] = get_base(
                                    cond)+self.__ir.imm
                        if cond & lt:
                            if self.__iregs[self.__priv][r_cond] < 0:
                                self.__sregs[self.__priv][Code] = get_base(
                                    cond)+self.__ir.imm
                    case 0x0b:  # call
                        r_cond, r_addr, cond = self.__ir.r1, self.__ir.r2, self.__ir.r3

                        def precall():
                            self.__tlb.update(self.__sregs[self.__priv])
                            self.__iregs[self.__priv][Stack] -= 8
                            self.__mar = int(self.__iregs[self.__priv][Stack])
                            self.__mdr = self.__iregs[self.__priv][RetAddr].view(
                                uint8)
                            self.__rwThread = Thread(target=self.__ram, daemon=True,
                                                     args=(8, Write | (Priv if self.__priv else 0)))
                            self.__rwThread.start()

                        def get_base(cond: int):
                            if cond & abs_addr:
                                return self.__iregs[self.__priv][r_addr]
                            else:
                                return self.__sregs[self.__priv][Code] & 0xffff_ffff_ffff_fc00

                        if cond & eq:
                            if self.__iregs[self.__priv][r_cond] == 0:
                                precall()
                                self.__iregs[self.__priv][RetAddr] = self.__sregs[self.__priv][Code]
                                self.__sregs[self.__priv][Code] = get_base(
                                    cond)+self.__ir.imm
                        if cond & gt:
                            if self.__iregs[self.__priv][r_cond] > 0:
                                precall()
                                self.__iregs[self.__priv][RetAddr] = self.__sregs[self.__priv][Code]
                                self.__sregs[self.__priv][Code] = get_base(
                                    cond)+self.__ir.imm
                        if cond & lt:
                            if self.__iregs[self.__priv][r_cond] < 0:
                                precall()
                                self.__iregs[self.__priv][RetAddr] = self.__sregs[self.__priv][Code]
                                self.__sregs[self.__priv][Code] = get_base(
                                    cond)+self.__ir.imm
                    case 0x0c:  # ret
                        r_cond, cond = self.__ir.r1, self.__ir.r3

                        def postcall():
                            self.__tlb.update(self.__sregs[self.__priv])
                            self.__mar = int(self.__iregs[self.__priv][Stack])
                            self.__rwThread = Thread(target=self.__ram, daemon=True,
                                                     args=(8, Read | (Priv if self.__priv else 0)))
                            self.__rwThread.start()
                            self.__iregs[self.__priv][Stack] += 8

                        if cond & eq:
                            if self.__iregs[self.__priv][r_cond] == 0:
                                postcall()
                                self.__sregs[self.__priv][Code] = self.__iregs[self.__priv][RetAddr]
                        if cond & gt:
                            if self.__iregs[self.__priv][r_cond] > 0:
                                postcall()
                                self.__sregs[self.__priv][Code] = self.__iregs[self.__priv][RetAddr]
                        if cond & lt:
                            if self.__iregs[self.__priv][r_cond] < 0:
                                postcall()
                                self.__sregs[self.__priv][Code] = self.__iregs[self.__priv][RetAddr]
                    case 0x0f:  # trap
                        raise Trap(self.__ir.imm)
                    case 0x20:  # mul
                        rd, rs, rt = self.__ir.r1, self.__ir.r2, self.__ir.r3
                        self.__iregs[self.__priv][rd] = self.__iregs[self.__priv][rs] * \
                            self.__iregs[self.__priv][rt]
                    case 0x21:  # div
                        rd, rs, rt = self.__ir.r1, self.__ir.r2, self.__ir.r3
                        if self.__iregs[self.__priv][rt] == 0:
                            raise DivideByZero(
                                self.__sregs[self.__priv][Code]-1)
                        self.__iregs[self.__priv][rd] = self.__iregs[self.__priv][rs] // \
                            self.__iregs[self.__priv][rt]
                    case 0x22:  # mod
                        rd, rs, rt = self.__ir.r1, self.__ir.r2, self.__ir.r3
                        if self.__iregs[self.__priv][rt] == 0:
                            raise DivideByZero(
                                self.__sregs[self.__priv][Code]-1)
                        self.__iregs[self.__priv][rd] = self.__iregs[self.__priv][rs] % \
                            self.__iregs[self.__priv][rt]
                    case 0x23:  # mulh
                        rd, rs, rt = self.__ir.r1, self.__ir.r2, self.__ir.r3
                        self.__iregs[self.__priv][rd] = (self.__iregs[self.__priv][rs] *
                                                         self.__iregs[self.__priv][rt]) >> 32
                    case 0x24:  # divh
                        rd, rs, rt = self.__ir.r1, self.__ir.r2, self.__ir.r3
                        if self.__iregs[self.__priv][rt] == 0:
                            raise DivideByZero(
                                self.__sregs[self.__priv][Code]-1)
                        self.__iregs[self.__priv][rd] = (self.__iregs[self.__priv][rs] //
                                                         self.__iregs[self.__priv][rt]) >> 32
                    case 0x70:  # iret
                        if self.__priv == 0:
                            raise PrivilegeViolation(
                                self.__sregs[self.__priv][Code]-1)
                        self.__priv = 0
                    case 0x71:  # kload
                        if self.__priv == 0:
                            raise PrivilegeViolation(
                                self.__sregs[self.__priv][Code]-1)
                        rd, rs = self.__ir.r1, self.__ir.r2
                        if rd == Code:
                            raise IllegalInstruction(
                                self.__sregs[self.__priv][Code]-1)
                        self.__sregs[True][rd] = self.__iregs[True][rs]
                    case 0x72:  # kstore
                        if self.__priv == 0:
                            raise PrivilegeViolation(
                                self.__sregs[self.__priv][Code]-1)
                        rd, rs = self.__ir.r1, self.__ir.r2
                        if rd == Code:
                            raise IllegalInstruction(
                                self.__sregs[self.__priv][Code]-1)
                        self.__iregs[True][rd] = self.__sregs[True][rs]
                    case 0x73:  # read
                        pass
                    case 0x74:  # write
                        pass
                    case 0x7f:  # halt
                        return
                    case _:
                        raise IllegalInstruction(
                            self.__sregs[self.__priv][Code]-1)
            except Exception as e:
                self.__post(e)
